import numpy as np
from LMdata.LMdataset import kp2idx, idx2kp, Skdict, idx2cate, cate2idx
from collections import OrderedDict

def findPeakPoint(heatmap):
    '''
    :param heatmap: shape of (N, 24, H, W)
    :return: idx np shape of (N,24,2)  x, y  %
    '''
    N,C,H,W = heatmap.shape
    ret = np.zeros((N,C,2))
    peaks_idxs = np.argmax(heatmap.reshape(N,C,-1), 2)
    idx_y, idx_x = np.unravel_index(peaks_idxs, dims=(H, W))
    ret[:, :, 0] = 1.0*idx_x/W
    ret[:, :, 1] = 1.0*idx_y/H
    return ret


def cal_NEscore(preds, labels, cates):
    '''
    :param preds: np.array of shape (N, num_keypoints, 2)  float or int, x y
    :param labels:  np.array of shape (N, num_keypoints, 3)  int, x y vis
    :param cates: np.array of shape (N, ) cate_idx, refer to
    :return:
    '''
    detail = {}
    N = preds.shape[0]

    # calculate Sk (N,1)
    Sk = np.ones((N,1),dtype=np.float32)
    for cate in Skdict:
        dx = labels[cates==cate,Skdict[cate][0], 0] - labels[cates==cate,Skdict[cate][1], 0]
        dy = labels[cates==cate,Skdict[cate][0], 1] - labels[cates==cate,Skdict[cate][1], 1]
        Sk[cates==cate,0] = np.sqrt(dx**2 + dy**2)
        # print 'Sk,',cate, Sk[(cates==cate)].min(), Sk[(cates==cate)].max()
    # cal Dk (N,24)
    Dk = np.sqrt((preds[:,:,0] - labels[:,:,0])**2 + (preds[:,:,1] - labels[:,:,1])**2)


    # cal masked relative error  (N,24)
    mask = (labels[:, :, 2] > 0)
    error = 1. * Dk / Sk * mask

    # cal by cate :
    for cate in Skdict:
        support = (cates==cate).sum()

        if support > 0:
            detail[idx2cate[cate]] = []
            if mask[cates==cate].sum()>0:
                detail[idx2cate[cate]].append(error[cates==cate].sum() / mask[cates==cate].sum())
            else:
                detail[idx2cate[cate]].append(0.)
            detail[idx2cate[cate]].append(support)
    return error.sum() / mask.sum(), detail
